package com.sunshine.bayes;

import java.util.List;
import java.util.Set;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;

public class Bayes {

    static final String[][] postingList = {{"my", "dog", "has", "flea", "problems", "help", "please"},
        {"maybe", "not", "take", "him", "to", "dog", "park", "stupid"},
        {"my", "dalmation", "is", "so", "cute", "I", "love", "him"},
        {"stop", "posting", "stupid", "worthless", "garbage"},
        {"mr", "licks", "ate", "my", "steak", "how", "to", "stop", "him"},
        {"quit", "buying", "worthless", "dog", "food", "stupid"}};//训练的数据集
    static final int[] classVec = {0, 1, 0, 1, 0, 1};//数据集的label，0是正面的，1是负面的

    //获取整个数据集的字典列表
    public static List<String> getDict() {
        Set<String> sets = Sets.newHashSet();
        for (String[] line : postingList) {
            for (String word : line) {
                sets.add(word);
            }
        }
        return Lists.newArrayList(sets);
    }

    //获取每一个训练数据集的字典向量列表
    public static int[] getClassVec(List<String> dicts, String[] line) {
        int[] vec = new int[dicts.size()];
        for (String word : line) {
            if (dicts.contains(word)) {
                vec[dicts.indexOf(word)] = 1;
            }
        }
        return vec;
    }

    //训练数据集，得到的结果是每一个词在正面和负面的概率是多少
    public static List<double[]> train(List<int[]> vec, int[] labels) {
        int[] p1 = new int[vec.get(0).length];
        int[] p2 = new int[vec.get(0).length];
        int p1Sum = 0;
        int p2Sum = 0;

        for (int i = 0; i < vec.size(); i++) {
            if (labels[i] == 1) {//负面的
                p1 = add(p1, vec.get(i));
                p1Sum += sum(vec.get(i));
            } else if (labels[i] == 0) {
                p2 = add(p2, vec.get(i));
                p2Sum += sum(vec.get(i));
            }

        }
        double[] p1Prob = div(p1, p1Sum);
        double[] p2Prob = div(p2, p2Sum);

        return Lists.newArrayList(p1Prob, p2Prob);
    }

    //预测给定的数据的正负面概率，0表示正面1表示负面
    public static double predict(String[] words, List<String> dicts, double[] p1Prob, double[] p2Prob) {
        int[] vec = getClassVec(dicts, words);

        double predictP1Prob = sum(mul(vec, p1Prob));
        double predictP2Prob = sum(mul(vec, p2Prob));
        //TODO 平滑技术
        if (predictP2Prob == 0) {
            return 1;
        }
        return predictP1Prob / predictP2Prob;
    }

    public static void main(String[] args) {
        List<String> dicts = getDict();
        List<int[]> trainVec = Lists.newArrayList();
        for (String[] line : postingList) {
            trainVec.add(getClassVec(dicts, line));
        }
        List<double[]> lists = train(trainVec, classVec);
        double[] p1Prob = lists.get(0);
        double[] p2Prob = lists.get(1);
        String[] test1 = {"love", "my", "dalmation"};
        String[] test2 = {"stupid", "garbage"};
        double testOne = predict(test1, dicts, p1Prob, p2Prob);
        double testTwo = predict(test2, dicts, p1Prob, p2Prob);
        System.out.println(testOne);
        System.out.println(testTwo);
    }

    private static double sum(double[] mul) {
        double v = 0;
        for (int i = 0; i < mul.length; i++) {
            v += mul[i];
        }
        return v;
    }

    private static int sum(int[] ints) {
        int v = 0;
        for (int i = 0; i < ints.length; i++) {
            v += ints[i];
        }
        return v;
    }

    private static double[] mul(int[] vec1, double[] vec2) {
        double[] v = new double[vec1.length];
        for (int i = 0; i < vec1.length; i++) {
            v[i] = vec1[i] * vec2[i];
        }
        return v;
    }

    private static double[] div(int[] vec, int total) {
        double[] v = new double[vec.length];

        for (int i = 0; i < vec.length; i++) {
            v[i] = vec[i] * 1.0 / total;
        }
        return v;
    }

    public static int[] add(int[] arr1, int[] arr2) {
        int[] v = new int[arr1.length];
        for (int i = 0; i < arr1.length; i++) {
            v[i] = arr1[i] + arr2[i];
        }
        return v;
    }
}